#---------------------------------------------------------------#
#    U-Net 
#        ver.1 - proto_UNet 用conv实现的最原始版本
#                有完整的网络框架，可在此版本上做网络结构改进优化
#---------------------------------------------------------------# 

import tensorflow as tf
import tensorflow.keras as keras

import numpy as np
from tensorflow.python.keras.backend import shape
from tensorflow.python.keras.layers.preprocessing.image_preprocessing import CenterCrop

'''
    
'''


def unet_conv_block(x, out_channels, kernel_size=(3,3), padding='valid', batch_norm_mode=False):
    y = keras.layers.Conv2D(out_channels, kernel_size, padding=padding, activation='relu', kernel_initializer='he_normal')(x)
    # if batch_norm_mode:
    #     y = keras.layers.BatchNormalization()(y) # 注意：BN要在Activation之前做
    # y = keras.layers.Activation('relu')(y)

    y = keras.layers.Conv2D(out_channels, kernel_size, padding=padding, activation='relu', kernel_initializer='he_normal')(y)
    # if batch_norm_mode:
    #     y = keras.layers.BatchNormalization()(y) # 注意：BN要在Activation之前做
    # y = keras.layers.Activation('relu')(y)

    return y

def unet_up_block(x1, x2, out_channels, kernel_size=(2,2), padding='valid', strides=2, up_mode='upconv'):
    '''
        x1 before upconv
        x2 encoder_map
    '''
    assert up_mode in ['upconv', 'upsample'], 'Wrong up_mode input: {}'.format(up_mode)

    if up_mode == 'upconv':
        y = keras.layers.Conv2DTranspose(out_channels, kernel_size, strides, padding)(x1)
    elif up_mode == 'upsample':
        y = keras.layers.UpSampling2D(size=kernel_size, interpolation='bilinear')(x1)
        y = keras.layers.Conv2D(out_channels, (1,1), padding='valid')(y)
    
    _, h1, w1, _ = y.shape
    _, h2, w2, _ = x2.shape
    dh, dw = int((h2-h1)/2), int((w2-w1)/2)
    x2 = x2[:, dh:h1+dh, dw:w1+dw, :]
    assert x2.shape[1:] == y.shape[1:], 'Mismatched shapes of x2 and y in decoder process!'

    y = tf.concat([y, x2], axis=-1)

    y = unet_conv_block(y, out_channels, (3,3), padding, False)
    
    return y

def f(photo):
    assert photo.shape[1:] == [572,572,3], 'Inappropriate input shape: {}'.format(photo.shape)

    channels = [64, 128, 256, 512, 1024]

    ### encoder
    o = photo
    encoder_maps = []
    for i in range(5):
        o = unet_conv_block(o, channels[i])
        if i < 4:
            encoder_maps.append(o)
            o = keras.layers.MaxPooling2D((2,2))(o) 

    assert o.shape[1:] == [28,28,1024], 'Wrong encoder process! '

    ### decoder
    for i in range(4):
        k = -1 * (i+1)
        o = unet_up_block(o, encoder_maps[k], channels[k-1], up_mode='upconv')

    return o

def get_net_model(net_input_shape, num_classes):
    # 1.输入层
    i_put = keras.layers.Input(net_input_shape)

    # 2.encoder + decoder
    decode_map = f(i_put)

    # 3.输出层
    o_put = keras.layers.Conv2D(num_classes+1, (1,1), padding='valid')(decode_map)
    o_put = keras.layers.Softmax()(o_put)

    # 4.模型实例
    net_model = keras.models.Model(i_put, o_put)

    return net_model

if __name__ == '__main__':
    # 测试搭建的网络
    net_input_shape = [572, 572, 3]
    num_classes     = 1
    net_model = get_net_model(net_input_shape, num_classes)
    # net_model.summary()

    net_input_shape.insert(0, 2)
    # x = np.random.random(net_input_shape).astype(np.float32) # 增加batch轴
    x = np.ones(net_input_shape).astype(np.float32) / 2
    y = net_model(x)
    y = y.numpy()
    print('测试网络的前向计算功能：')
    print('    x.shape = {}'.format(x.shape))
    print('    y.shape = {}'.format(y.shape))